/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.starter.dubbo;

import com.alibaba.dubbo.common.Constants;
import com.alibaba.dubbo.common.extension.Activate;
import com.alibaba.dubbo.remoting.TimeoutException;
import com.alibaba.dubbo.rpc.Filter;
import com.alibaba.dubbo.rpc.Invocation;
import com.alibaba.dubbo.rpc.Invoker;
import com.alibaba.dubbo.rpc.Result;
import com.alibaba.dubbo.rpc.RpcContext;
import com.alibaba.dubbo.rpc.RpcException;
import com.alibaba.dubbo.rpc.support.RpcUtils;
import com.lvyh.lightframe.tracer.core.LogData;
import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.TraceContants;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContext;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContextHolder;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import com.lvyh.lightframe.tracer.core.samplers.TraceSamplerConfig;
import com.lvyh.lightframe.tracer.core.util.StringUtils;
import com.lvyh.lightframe.tracer.plugins.dubbo.AttachmentKeyConstants;
import com.lvyh.lightframe.tracer.plugins.dubbo.tracer.DubboConsumerTracer;
import com.lvyh.lightframe.tracer.starter.util.SpringContextUtils;
import io.opentracing.tag.Tags;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.CompletableFuture;

@Activate(group = {Constants.CONSUMER})
public class DubboConsumerTracerFilter implements Filter {
    private static final Logger logger = LoggerFactory.getLogger(DubboConsumerTracerFilter.class);

    private AbstractTracer dubboConsumerTracer;

    public DubboConsumerTracerFilter() {
        TraceSamplerConfig traceSamplerConfig = SpringContextUtils.getBean(TraceSamplerConfig.class);
        AbstractTracer dubboConsumerTracer = DubboConsumerTracer.getDubboConsumerTracer(traceSamplerConfig);
        this.dubboConsumerTracer = dubboConsumerTracer;
    }

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
        logger.info("[DubboConsumerTracerFilter] entry in consumer trace filter");

        if ("$echo".equals(invocation.getMethodName())) {
            return invoker.invoke(invocation);
        }

        Result result = doFilter(invoker, invocation);
        return result;
    }

    private Result doFilter(Invoker<?> invoker, Invocation invocation) {
        RpcContext rpcContext = RpcContext.getContext();
        String methodName = rpcContext.getMethodName();
        String service = invoker.getInterface().getSimpleName();

        String operationName = service + "#" + methodName;
        SimpleSpan simpleSpan = dubboConsumerTracer.clientSend(operationName);

        setRpcClientSpanTag(invoker, simpleSpan);

        String serializeValue = simpleSpan.getSimpleSpanContext().serialize();
        invocation.getAttachments().put(TraceContants.RPC_TRACE_NAME, serializeValue);

        // Enable asynchronous
        boolean isAsync = RpcUtils.isAsync(invoker.getUrl(), invocation);

        // Is it a one-way send
        boolean isOneway = RpcUtils.isOneway(invoker.getUrl(), invocation);

        if (isAsync) {
            simpleSpan.setTag(TraceContants.RPC_INVOKE_TYPE, "future");
        } else {
            if (isOneway) {
                simpleSpan.setTag(TraceContants.RPC_INVOKE_TYPE, "oneway");
            } else {
                simpleSpan.setTag(TraceContants.RPC_INVOKE_TYPE, "sync");
            }
        }

        Result result;
        Throwable exception = null;
        String resultCode = TraceContants.RESULT_CODE_SUCCESS;
        try {
            result = invoker.invoke(invocation);
            if (result == null) {
                if (isOneway) {
                    simpleSpan.setTag(TraceContants.RESP_SIZE, 0);
                }
            } else {
                appendElapsedTimeTags(invocation, simpleSpan, result, true);
            }
        } catch (RpcException e) {
            exception = e;
            throw e;
        } catch (Throwable t) {
            exception = t;
            throw new RpcException(t);
        } finally {
            if (exception != null) {
                if (exception instanceof RpcException) {
                    simpleSpan.setTag(Tags.ERROR.getKey(), exception.getMessage());
                    RpcException rpcException = (RpcException) exception;
                    resultCode = String.valueOf(rpcException.getCode());
                } else {
                    resultCode = TraceContants.RESULT_CODE_ERROR;
                }
            }

            SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
            SimpleSpan currentSpan = simpleTraceContext.getCurrentSpan();
            Sampler sampler = dubboConsumerTracer.getSimpleTracer().getSampler();
            //sampler is support &  current span is root span
            if (sampler != null && currentSpan.getParentSpan() == null) {
                currentSpan.getSimpleSpanContext().setSampled(true);
                currentSpan.getSimpleSpanContext().setSampler(sampler);
            }

            if (!isAsync) {
                dubboConsumerTracer.clientReceive(resultCode);
            } else {
                SimpleTraceContext sofaTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
                SimpleSpan clientSpan = sofaTraceContext.pop();
                if (clientSpan != null) {
                    simpleSpan.log(LogData.CLIENT_SEND_EVENT_VALUE);
                }

                if (clientSpan != null && clientSpan.getParentSpan() != null) {
                    sofaTraceContext.add(clientSpan.getParentSpan());
                }

                CompletableFuture<Object> future = (CompletableFuture<Object>) RpcContext.getContext().getFuture();
                future.whenComplete((object, throwable) -> {
                    if (throwable instanceof TimeoutException) {
                        clientSpan.setTag(Tags.ERROR.getKey(), throwable.getMessage());
                        dubboConsumerTracer.clientReceiveTagFinish(clientSpan, TraceContants.RESULT_CODE_TIME_OUT);
                    }
                });
            }
        }
        return result;
    }


    private static void appendElapsedTimeTags(Invocation invocation, SimpleSpan simpleSpan, Result result, boolean isClient) {
        String reqSize;
        String respSize;
        String elapsed;
        String deElapsed;
        if (isClient) {
            reqSize = invocation.getAttachment(AttachmentKeyConstants.CLIENT_SERIALIZE_SIZE);
            elapsed = invocation.getAttachment(AttachmentKeyConstants.CLIENT_SERIALIZE_TIME);
            respSize = result.getAttachment(AttachmentKeyConstants.CLIENT_DESERIALIZE_SIZE);
            deElapsed = result.getAttachment(AttachmentKeyConstants.CLIENT_DESERIALIZE_TIME);

            simpleSpan.setTag(AttachmentKeyConstants.CLIENT_SERIALIZE_TIME,
                    parseAttachment(elapsed, 0));
            simpleSpan.setTag(AttachmentKeyConstants.CLIENT_DESERIALIZE_TIME,
                    parseAttachment(deElapsed, 0));
            simpleSpan.setTag(AttachmentKeyConstants.CLIENT_SERIALIZE_SIZE,
                    parseAttachment(reqSize, 0));
            simpleSpan.setTag(AttachmentKeyConstants.CLIENT_DESERIALIZE_SIZE,
                    parseAttachment(respSize, 0));
        } else {
            reqSize = invocation.getAttachment(AttachmentKeyConstants.SERVER_DESERIALIZE_SIZE);
            deElapsed = invocation.getAttachment(AttachmentKeyConstants.SERVER_DESERIALIZE_TIME);
            respSize = result.getAttachment(AttachmentKeyConstants.SERVER_SERIALIZE_SIZE);
            elapsed = result.getAttachment(AttachmentKeyConstants.SERVER_SERIALIZE_TIME);

            simpleSpan.setTag(AttachmentKeyConstants.SERVER_DESERIALIZE_SIZE,
                    parseAttachment(reqSize, 0));
            simpleSpan.setTag(AttachmentKeyConstants.SERVER_DESERIALIZE_TIME,
                    parseAttachment(deElapsed, 0));
            simpleSpan.setTag(AttachmentKeyConstants.SERVER_SERIALIZE_SIZE,
                    parseAttachment(respSize, 0));
            simpleSpan.setTag(AttachmentKeyConstants.SERVER_SERIALIZE_TIME,
                    parseAttachment(elapsed, 0));
        }

    }

    private static int parseAttachment(String value, int defaultVal) {
        try {
            if (StringUtils.isNotBlank(value)) {
                defaultVal = Integer.parseInt(value);
            }
        } catch (Exception e) {
            logger.error("Failed to parse Dubbo plugin params.", e);
        }
        return defaultVal;
    }

    private void setRpcClientSpanTag(Invoker<?> invoker, SimpleSpan simpleSpan) {
        RpcContext rpcContext = RpcContext.getContext();
        Map<String, String> tags = simpleSpan.getTags();

        tags.put(Tags.SPAN_KIND.getKey(), rpcContext.isConsumerSide() ? Tags.SPAN_KIND_CLIENT : Tags.SPAN_KIND_SERVER);

        String protocol = rpcContext.getUrl().getProtocol();
        tags.put(TraceContants.PROTOCOL, protocol == null ? "" : protocol);

        String service = invoker.getInterface().getName();
        tags.put(TraceContants.SERVICE, service == null ? "" : service);

        String methodName = rpcContext.getMethodName();
        tags.put(TraceContants.METHOD, methodName == null ? "" : methodName);

        tags.put(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());

        String app = rpcContext.getUrl().getParameter(Constants.APPLICATION_KEY);
        tags.put(TraceContants.LOCAL_APP, app);
        tags.put(TraceContants.REMOTE_APP, app + Tags.SPAN_KIND_CONSUMER);

        tags.put(TraceContants.REMOTE_HOST, rpcContext.getRemoteHost());
        tags.put(TraceContants.REMOTE_PORT, String.valueOf(rpcContext.getRemotePort()));
        tags.put(TraceContants.LOCAL_HOST, rpcContext.getLocalHost());
    }

}